from numpy import array, where

"""
Align takes 5 arguments:
query is the query sequence (generally longer than reference)
reference is the reference sequence (generally shorter)
t is a numpy array
    must be initialized to numrows = len(longest query)+1 
                         x numcols = len(reference)+1
    first column and row should be sequence increasing by 1
index is the last row processed in the previous sequence
dist is the max edit dist

shape of the array is

         reference
         -ACTTG
         012345
    q - 0  
    u A 1
    e G 2
    r T 3
    y C 4
      C 5

"""

def align(query,reference,t,index,dist):
    

    #===============================================================================================
    # First, get the lengths of each string because we only need to compute that much
    # of the DP table
    #===============================================================================================
    lenq = len(query)
    lenr = len(reference)
    
    #===============================================================================================
    # Each time we call align, we need to clear out the fresh part of the table (below our index)
    #===============================================================================================
    t[-(t.shape[0]-index):,1:].fill(dist + 1)
    
    #===============================================================================================
    # Starting at the row indicated by index, we start moving row by row, down the DP table
    # Remember that index will usually be 1 + LCP(thisQuery, lastQuery)
    # This allows us to reuse the work we did the last alignment
    #===============================================================================================
    for row in range(index,lenq + 1):
        #===========================================================================================
        # Where() returns the indeces (columns) which match a condition
        # In this case we want all columns of the previous row still within the edit distance
        #===========================================================================================
        colrange, = where(t[row-1] <= dist)
        #===========================================================================================
        # If no values in the previous row were less than the edit distance,
        # return the row 2 previous (which would be the last valid row)
        #===========================================================================================
        if colrange.size == 0:
            break
        #===========================================================================================
        # Increase all column indeces by 1 to move diagonally down the table
        #===========================================================================================
        colrange += 1
        #===========================================================================================
        # If moving diagonally, moved us out of bounds we need to make sure we still have a valid
        # index to look at
        #===========================================================================================
        if colrange.max() > lenr:
            #=======================================================================================
            # No valid index, return the row 2 previous... same as if there had been no matches
            # this row
            #=======================================================================================
            if colrange.min() > lenr: break 
            #=======================================================================================
            # Even though some of the indeces were out of bounds, we have some good ones still
            # so we'll truncate colrange to only include those
            #=======================================================================================
            colrange = colrange[colrange<=lenr]
        #===========================================================================================
        # Now, we can look at the range of indeces we just found and use our DP function
        #===========================================================================================
        for col in range(colrange.min(),colrange.max()+1):
            if reference[col-1] == query[row-1]: mismatch = 0 #if they match, no penalty
            else: mismatch = 1 #else penalty
            #=======================================================================================
            # Standard dynamic programming function
            #=======================================================================================
            f = array([t[row-1,col]+1,
                       t[row,col-1]+1,
                       t[row-1,col-1]+mismatch])
            t[row,col] = f.min()
        
    #===============================================================================================
    # Here is where truncating the sequences comes in handy...
    # If there is a valid alignment along the right edge of the table its a good alignment
    #===============================================================================================
    if t[:lenq,lenr].min() <= dist:
        return True, row - 1
    #===============================================================================================
    # We might've missed alignments of shorter queries... so check if the query is shorter
    #===============================================================================================
    elif (lenr > lenq): 
        #===========================================================================================
        # If it is, then check for an alignment along the bottom row
        #===========================================================================================
        if (t[lenq].min() <= dist):
            return True, lenq
        else:
            #=======================================================================================
            # Row - 1 is the last valid row because row never advances past lenq
            #=======================================================================================
            return False, row - 1
    else:
        return False, row - 2
